import torch
# data
TRAIN_BATCH_SIZE = 64
TEST_BATCH_SIZE = 16

# model
# FCN
# MODEL_NAME = 'FCN'
INPUT_SIZE = 784
HIDDEN_SIZE1 = 256
HIDDEN_SIZE2 = 512
OUTPUT_SIZE =  10
# CNN
MODEL_NAME = 'CNN'

# train
BATCH_SIZE = 16
DEVICE = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
EPOCHS = 20
LEARNING_RATE = 0.01